from sparkai.core.outputs import LLMResult
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

from config_set import config_instance

SPARKAI_URL = config_instance.get('SPARKAI_URL')
SPARKAI_APP_ID = config_instance.get('SPARKAI_APP_ID')
SPARKAI_API_SECRET = config_instance.get('SPARKAI_API_SECRET')
SPARKAI_API_KEY = config_instance.get('SPARKAI_API_KEY')

#星火认知大模型Spark Max的domain值，其他版本大模型domain值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_DOMAIN = '4.0Ultra'

def get_xinghuo_response(user_input):
    """
    获取星火大模型的回复
    :param user_input: 用户输入的内容
    :return: 大模型的回复
    """
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=True,
        request_timeout=10
    )
    messages = [ChatMessage(
        role="user",
        content=user_input
    )]
    handler = ChunkPrintHandler()
    response = spark.generate([messages], callbacks=[handler])
    return response


if __name__ == '__main__':
    res = get_xinghuo_response("你好啊！")
    # 从结果中提取回复文本
    if isinstance(res, LLMResult):
        # 通常 generations 是一个二维列表，这里假设取第一个生成结果
        first_generation = res.generations[0][0]
        reply_text = first_generation.text
        # print(reply_text)
    else:
        print("结果不是 LLMResult 类型")


